#include <iostream>
#include <unistd.h>
#include <thread>
#include <vector>
#include <memory>

#include <cuda_runtime.h>

#define CHECK_CUDA_ERROR(call) \
    do { \
        cudaError_t err = (call); \
        if (err != cudaSuccess) { \
            std::cerr << "CUDA error at " << __FILE__ << ":" << __LINE__ << ": " \
                      << cudaGetErrorString(err) << std::endl; \
        } \
    } while(0)

const char* cudaStreamCaptureModeToString(cudaStreamCaptureMode mode) {
    switch (mode) {
        case cudaStreamCaptureModeGlobal:
            return "cudaStreamCaptureModeGlobal";
        case cudaStreamCaptureModeThreadLocal:
            return "cudaStreamCaptureModeThreadLocal";
        case cudaStreamCaptureModeRelaxed:
            return "cudaStreamCaptureModeRelaxed";
        default:
            return "Unknown cudaStreamCaptureMode";
    }
}

//ref paddle
struct CUDAGraphCaptureModeGuard {
    explicit CUDAGraphCaptureModeGuard(cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed) {
        CHECK_CUDA_ERROR(cudaThreadExchangeStreamCaptureMode(&mode));
        std::cout <<"-----[GUARD] begin" << std::endl;
        old_mode_ = mode;
    }

    ~CUDAGraphCaptureModeGuard() {
        std::cout <<"-----[GUARD] end" << std::endl;
        CHECK_CUDA_ERROR(cudaThreadExchangeStreamCaptureMode(&old_mode_));
    }
    cudaStreamCaptureMode old_mode_;
};

void func_capture(int t_index){
    id_t tid;
    std::shared_ptr<CUDAGraphCaptureModeGuard> guard;
    if(t_index == 4)
        guard = std::make_shared<CUDAGraphCaptureModeGuard>();

    for(int i = 0; i < (t_index+1)*10000000; ++i){
        tid = gettid();
    }
    std::cout << t_index << " tid " << tid << std::endl;
    void* d_data;
    CHECK_CUDA_ERROR(cudaMalloc(&d_data, 1024));
    CHECK_CUDA_ERROR(cudaFree(d_data));
}

int main() {
    std::vector<std::thread> threads;

    id_t tid = gettid();
    std::cout << "_ tid " << tid << std::endl;

    cudaStream_t stream;
    CHECK_CUDA_ERROR(cudaStreamCreate(&stream));    
    CHECK_CUDA_ERROR(cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal));

    for (int t_index = 0; t_index < 5; ++t_index) {
        threads.emplace_back([t_index]() {
            func_capture(t_index);
        });
    }
    for (auto& t : threads) {
        t.join();
    }

    cudaGraph_t graph;
    CHECK_CUDA_ERROR(cudaStreamEndCapture(stream, &graph));
    
    CHECK_CUDA_ERROR(cudaStreamDestroy(stream));
    CHECK_CUDA_ERROR(cudaGraphDestroy(graph)); 

    return 0;
}
